package com.kara.woodAgent.agent.graph.impl;

import com.kara.woodAgent.agent.graph.Next;
import com.kara.woodAgent.agent.graph.Node;
import com.kara.woodAgent.agent.graph.communicate.ChannelMessage;
import com.kara.woodAgent.agent.model.ModelProvider;
import com.kara.woodAgent.agent.tool.provider.ToolProvider;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.response.ChatResponse;

import java.util.List;

/**
 * @author tzy
 * @version 1.0
 * @project wood_agent
 * @description
 * @date 2025/6/4 11:30:46
 */
public class ModelCheckpointNode implements Node<ReactGraphContext> {

    private final ModelProvider modelProvider;

    private final ToolProvider.ToolServiceHolder toolProviderHolder;


    public ModelCheckpointNode(ModelProvider modelProvider, ToolProvider.ToolServiceHolder toolProviderHolder) {
        this.modelProvider = modelProvider;
        this.toolProviderHolder = toolProviderHolder;
    }

    @Override
    public Next execute(ReactGraphContext context) {
        List<ChatMessage> chatMessages = context.getChatMessages();
        if (!context.getChannel().checking()) {

            ChatLanguageModel chatModel = modelProvider.getChatModel();

            ChatRequest request = ChatRequest.builder()
                    .toolSpecifications(toolProviderHolder.toolSpecifications())
                    .messages(chatMessages).build();
            ChatResponse chatResponse = chatModel.chat(request);

            AiMessage aiMessage = chatResponse.aiMessage();
            chatMessages.add(aiMessage);
            ModelCheckpointPayload payload = ModelCheckpointPayload.of(aiMessage.text());
            context.getChannel().checkpoint(payload);
        }

        ChannelMessage message = context.getChannel().resume();
        String content = message.getContent();
        if ("y".equals(content)) {
            AiMessage aiMessage = (AiMessage) chatMessages.getLast();

            if (aiMessage.hasToolExecutionRequests()) {
                context.setToolExecutionRequest(aiMessage.toolExecutionRequests());
                return Next.NextNode("tool").step(aiMessage.text() + " " + aiMessage.toolExecutionRequests().toString());
            } else {
                return Next.End().step(aiMessage.text());
            }
        } else {
            UserMessage userMessage = UserMessage.from("不同意你对于这个步骤的想法,请重新思考给出你的建议");
            chatMessages.add(userMessage);
            return Next.NextNode("model").step(userMessage.toString());
        }

    }
}
